# A marching cube test
import numpy as np

try:
    import numba as nb
    numba_present = True
except ImportError:
    import warnings
    warnings.warn(
    "Numba not installed. The isosurface generation will be extremely slow",
        ImportWarning)
    numba_present = False


def isosurface_from_function(function, extents=[[-1, 1], [-1, 1], [-1, 1]], 
                                       isolevel=0.3, resolution=32):
    area_min, area_max = np.array(extents).T
    spacing = (area_max - area_min)/resolution
    
    x = np.linspace(area_min[0], area_max[0], resolution)
    y = np.linspace(area_min[1], area_max[1], resolution)
    z = np.linspace(area_min[2], area_max[2], resolution)
    xv, yv, zv = np.meshgrid(x, y, z)
    
    return isosurface_from_data(function(xv, yv, zv), isolevel, 
                                  area_min, spacing)

def isosurface_from_data(data, isolevel, origin, spacing):
    """Small wrapper to get directly vertices and faces to feed into programs
    """
    spacing = np.array(extent/resolution)
    if isolevel >= 0:
        triangles = marching_cubes(data, isolevel)
    else: # Wrong traingle unwinding roder -- god only knows why
        triangles = marching_cubes(-data, -isolevel)
    faces = []
    verts = []
    for i, t in enumerate(triangles):
       faces.append([i * 3, i * 3 +1, i * 3 + 2])
       verts.extend(t)
       
    faces = np.array(faces)
    verts = origin + spacing/2 + np.array(verts)*spacing
    return verts, faces

def marching_cubes(field, isolevel):
    # The field is like gridpoints, and gridpoints define cubes.
    triangles = []
    # Here we determine the value for each field
    for i in range(field.shape[0] - 1):
        for j in range(field.shape[1] - 1):
            for k in range(field.shape[2] - 1):

                points = {
                    0: (i, j, k),
                    1: (i, j+1, k),
                    2: (i+1, j+1, k),
                    3: (i+1, j, k),
                    4: (i, j, k+1),
                    5: (i, j+1, k+1),
                    6: (i+1, j+1, k+1),
                    7: (i+1, j, k+1)
                }
                # Which kind of cube is this guy?
                cube_index = 0
                if field[points[0]] < isolevel: cube_index |= 1
                if field[points[1]] < isolevel: cube_index |= 2
                if field[points[2]] < isolevel: cube_index |= 4
                if field[points[3]] < isolevel: cube_index |= 8
                if field[points[4]] < isolevel: cube_index |= 16
                if field[points[5]] < isolevel: cube_index |= 32
                if field[points[6]] < isolevel: cube_index |= 64
                if field[points[7]] < isolevel: cube_index |= 128

                # Get the faces from the cube

                for e1,e2,e3 in tris_as_edges[cube_index]:
                    # Get the three interpolated points on the edges
                    triangle = []
                    for edge in (e1, e2, e3):
                        s, e = edge2pts[edge]
                        p1, p2 = np.array(points[e], 'f'), np.array(points[s], 'f')
                        v = interpolate_edge_coordinates(p1, field[points[e]],
                                                         p2, field[points[s]],
                                                         isolevel)
                        triangle.append(v)
                    triangles.append(triangle)
    triangles_ = np.array(triangles) # triangles_ NUMBA BUG

    if len(triangles) == 0:
        return np.array([])

    # TODO Let's just invert for now, but no one knows what the problem is
    triangles_[:, :, [0, 1]] = triangles_[:, :, [1, 0]]
    return triangles_



def interpolate_edge_coordinates(point1, value1, point2, value2, isolevel):
    return point1 + (isolevel - value1) * (point2 - point1)/(value2 - value1)

if numba_present:
    # marching_cubes = nb.jit('pyobject(f4[:, :, :], f4)')(marching_cubes)
    # marching_cubes = nb.jit()(marching_cubes)
    interpolate_edge_coordinates = nb.jit(interpolate_edge_coordinates)

# Returns the vertices that make up an edge
edge2pts = [(0,1),(1,2),(2,3),(3,0),(4,5),(5,6),(6,7),(7,4),
            (0,4),(1,5),(2,6),(3,7)]

# For a given case (0-255), this array yields the corresponding
# triangles, given as a tuple of three edges
tris_as_edges = [
    [],
    [(0, 8, 3)],
    [(0, 1, 9)],
    [(1, 8, 3), (9, 8, 1)],
    [(1, 2, 10)],
    [(0, 8, 3), (1, 2, 10)],
    [(10, 9, 0), (0, 2, 10)],
    [(2, 8, 3), (2, 10, 8), (10, 9, 8)],
    [(3, 11, 2)],
    [(2, 0, 8), (8, 11, 2)],
    [(1, 9, 0), (2, 3, 11)],
    [(1, 11, 2), (1, 9, 11), (9, 8, 11)],
    [(3, 11, 1), (1, 11, 10)],
    [(0, 10, 1), (0, 8, 10), (11, 10, 8)],
    [(3, 9, 0), (3, 11, 9), (11, 10, 9)],
    [(9, 8, 11), (10, 9, 11)],
    [(4, 7, 8)],
    [(4, 3, 0), (7, 3, 4)],
    [(0, 1, 9), (8, 4, 7)],
    [(4, 1, 9), (4, 7, 1), (7, 3, 1)],
    [(1, 2, 10), (8, 4, 7)],
    [(3, 4, 7), (3, 0, 4), (1, 2, 10)],
    [(9, 2, 10), (9, 0, 2), (8, 4, 7)],
    [(2, 10, 9), (2, 9, 7), (2, 7, 3), (7, 9, 4)],
    [(8, 4, 7), (3, 11, 2)],
    [(11, 4, 7), (11, 2, 4), (2, 0, 4)],
    [(9, 0, 1), (8, 4, 7), (2, 3, 11)],
    [(4, 7, 11), (9, 4, 11), (9, 11, 2), (2, 1, 9)],
    [(3, 11, 10), (10, 1, 3), (7, 8, 4)],
    [(11, 10, 1), (4, 11, 1), (7, 11, 4), (4, 1, 0)],
    [(4, 7, 8), (3, 11, 9), (9, 11, 10), (9, 0, 3)],
    [(4, 7, 11), (4, 11, 9), (9, 11, 10)],
    [(9, 5, 4)],
    [(9, 5, 4), (0, 8, 3)],
    [(0, 5, 4), (1, 5, 0)],
    [(8, 5, 4), (8, 3, 5), (3, 1, 5)],
    [(1, 2, 10), (9, 5, 4)],
    [(3, 0, 8), (1, 2, 10), (4, 9, 5)],
    [(5, 2, 10), (5, 4, 2), (4, 0, 2)],
    [(2, 10, 5), (3, 2, 5), (3, 5, 4), (3, 4, 8)],
    [(9, 5, 4), (2, 3, 11)],
    [(0, 11, 2), (0, 8, 11), (4, 9, 5)],
    [(0, 5, 4), (0, 1, 5), (2, 3, 11)],
    [(2, 1, 5), (2, 5, 8), (2, 8, 11), (4, 8, 5)],
    [(10, 3, 11), (10, 1, 3), (9, 5, 4)],
    [(4, 9, 5), (0, 8, 1), (8, 10, 1), (8, 11, 10)],
    [(5, 4, 0), (5, 0, 11), (5, 11, 10), (11, 0, 3)],
    [(5, 4, 8), (5, 8, 10), (10, 8, 11)],
    [(9, 7, 8), (5, 7, 9)],
    [(9, 3, 0), (9, 5, 3), (5, 7, 3)],
    [(0, 7, 8), (0, 1, 7), (1, 5, 7)],
    [(1, 5, 3), (3, 5, 7)],
    [(9, 7, 8), (9, 5, 7), (10, 1, 2)],
    [(10, 1, 2), (9, 5, 0), (5, 3, 0), (5, 7, 3)],
    [(8, 0, 2), (8, 2, 5), (8, 5, 7), (10, 5, 2)],
    [(2, 10, 5), (2, 5, 3), (3, 5, 7)],
    [(7, 9, 5), (7, 8, 9), (3, 11, 2)],
    [(9, 5, 7), (9, 7, 2), (9, 2, 0), (2, 7, 11)],
    [(2, 3, 11), (0, 1, 8), (1, 7, 8), (1, 5, 7)],
    [(11, 2, 1), (11, 1, 7), (7, 1, 5)],
    [(9, 5, 8), (8, 5, 7), (10, 1, 3), (10, 3, 11)],
    [(5, 7, 0), (5, 0, 9), (7, 11, 0), (1, 0, 10), (11, 10, 0)],
    [(11, 10, 0), (11, 0, 3), (10, 5, 0), (8, 0, 7), (5, 7, 0)],
    [(5, 11, 10), (11, 5, 7)],
    [(10, 6, 5)],
    [(0, 8, 3), (6, 5, 10)],
    [(9, 0, 1), (5, 10, 6)],
    [(1, 8, 3), (1, 9, 8), (5, 10, 6)],
    [(1, 6, 5), (2, 6, 1)],
    [(1, 6, 5), (1, 2, 6), (3, 0, 8)],
    [(9, 6, 5), (9, 0, 6), (0, 2, 6)],
    [(5, 9, 8), (5, 8, 2), (5, 2, 6), (3, 2, 8)],
    [(2, 3, 11), (10, 6, 5)],
    [(11, 0, 8), (11, 2, 0), (10, 6, 5)],
    [(0, 1, 9), (2, 3, 11), (5, 10, 6)],
    [(5, 10, 6), (1, 9, 2), (9, 11, 2), (9, 8, 11)],
    [(6, 3, 11), (6, 5, 3), (5, 1, 3)],
    [(0, 8, 11), (0, 11, 5), (0, 5, 1), (5, 11, 6)],
    [(3, 11, 6), (0, 3, 6), (0, 6, 5), (0, 5, 9)],
    [(6, 5, 9), (6, 9, 11), (11, 9, 8)],
    [(5, 10, 6), (4, 7, 8)],
    [(4, 3, 0), (4, 7, 3), (6, 5, 10)],
    [(1, 9, 0), (5, 10, 6), (8, 4, 7)],
    [(10, 6, 5), (1, 9, 7), (1, 7, 3), (7, 9, 4)],
    [(6, 1, 2), (6, 5, 1), (4, 7, 8)],
    [(1, 2, 5), (5, 2, 6), (3, 0, 4), (3, 4, 7)],
    [(8, 4, 7), (9, 0, 5), (0, 6, 5), (0, 2, 6)],
    [(7, 3, 9), (7, 9, 4), (3, 2, 9), (5, 9, 6), (2, 6, 9)],
    [(3, 11, 2), (7, 8, 4), (10, 6, 5)],
    [(5, 10, 6), (4, 7, 2), (4, 2, 0), (2, 7, 11)],
    [(0, 1, 9), (4, 7, 8), (2, 3, 11), (5, 10, 6)],
    [(9, 2, 1), (9, 11, 2), (9, 4, 11), (7, 11, 4), (5, 10, 6)],
    [(8, 4, 7), (3, 11, 5), (3, 5, 1), (5, 11, 6)],
    [(5, 1, 11), (5, 11, 6), (1, 0, 11), (7, 11, 4), (0, 4, 11)],
    [(0, 5, 9), (0, 6, 5), (0, 3, 6), (11, 6, 3), (8, 4, 7)],
    [(6, 5, 9), (6, 9, 11), (4, 7, 9), (7, 11, 9)],
    [(6, 4, 9), (9, 10, 6)],
    [(4, 10, 6), (4, 9, 10), (0, 8, 3)],
    [(10, 0, 1), (10, 6, 0), (6, 4, 0)],
    [(8, 3, 1), (8, 1, 6), (8, 6, 4), (6, 1, 10)],
    [(1, 4, 9), (1, 2, 4), (2, 6, 4)],
    [(3, 0, 8), (1, 2, 9), (2, 4, 9), (2, 6, 4)],
    [(0, 2, 4), (4, 2, 6)],
    [(8, 3, 2), (8, 2, 4), (4, 2, 6)],
    [(10, 4, 9), (10, 6, 4), (11, 2, 3)],
    [(0, 8, 2), (2, 8, 11), (4, 9, 10), (4, 10, 6)],
    [(3, 11, 2), (0, 1, 6), (0, 6, 4), (6, 1, 10)],
    [(6, 4, 1), (6, 1, 10), (4, 8, 1), (2, 1, 11), (8, 11, 1)],
    [(9, 6, 4), (9, 3, 6), (9, 1, 3), (11, 6, 3)],
    [(8, 11, 1), (8, 1, 0), (11, 6, 1), (9, 1, 4), (6, 4, 1)],
    [(3, 11, 6), (3, 6, 0), (0, 6, 4)],
    [(6, 4, 8), (8, 11, 6)],
    [(7, 10, 6), (7, 8, 10), (8, 9, 10)],
    [(0, 7, 3), (0, 10, 7), (0, 9, 10), (6, 7, 10)],
    [(10, 6, 7), (1, 10, 7), (1, 7, 8), (1, 8, 0)],
    [(10, 6, 7), (10, 7, 1), (1, 7, 3)],
    [(1, 2, 6), (1, 6, 8), (1, 8, 9), (8, 6, 7)],
    [(2, 6, 9), (2, 9, 1), (6, 7, 9), (0, 9, 3), (7, 3, 9)],
    [(7, 8, 0), (7, 0, 6), (6, 0, 2)],
    [(7, 3, 2), (6, 7, 2)],
    [(2, 3, 11), (10, 6, 8), (10, 8, 9), (8, 6, 7)],
    [(2, 0, 7), (2, 7, 11), (0, 9, 7), (6, 7, 10), (9, 10, 7)],
    [(1, 8, 0), (1, 7, 8), (1, 10, 7), (6, 7, 10), (2, 3, 11)],
    [(11, 2, 1), (11, 1, 7), (10, 6, 1), (6, 7, 1)],
    [(8, 9, 6), (8, 6, 7), (9, 1, 6), (11, 6, 3), (1, 3, 6)],
    [(0, 9, 1), (11, 6, 7)],
    [(7, 8, 0), (7, 0, 6), (3, 11, 0), (11, 6, 0)],
    [(7, 11, 6)],
    [(7, 6, 11)],
    [(3, 0, 8), (11, 7, 6)],
    [(0, 1, 9), (11, 7, 6)],
    [(8, 1, 9), (8, 3, 1), (11, 7, 6)],
    [(10, 1, 2), (6, 11, 7)],
    [(1, 2, 10), (3, 0, 8), (6, 11, 7)],
    [(2, 9, 0), (2, 10, 9), (6, 11, 7)],
    [(6, 11, 7), (2, 10, 3), (10, 8, 3), (10, 9, 8)],
    [(7, 2, 3), (6, 2, 7)],
    [(7, 0, 8), (7, 6, 0), (6, 2, 0)],
    [(2, 7, 6), (2, 3, 7), (0, 1, 9)],
    [(1, 6, 2), (1, 8, 6), (1, 9, 8), (8, 7, 6)],
    [(10, 7, 6), (10, 1, 7), (1, 3, 7)],
    [(10, 7, 6), (1, 7, 10), (1, 8, 7), (1, 0, 8)],
    [(0, 3, 7), (0, 7, 10), (0, 10, 9), (6, 10, 7)],
    [(7, 6, 10), (7, 10, 8), (8, 10, 9)],
    [(8, 4, 6), (6, 11, 8)],
    [(3, 6, 11), (3, 0, 6), (0, 4, 6)],
    [(8, 6, 11), (8, 4, 6), (9, 0, 1)],
    [(9, 4, 6), (9, 6, 3), (9, 3, 1), (11, 3, 6)],
    [(6, 8, 4), (6, 11, 8), (2, 10, 1)],
    [(1, 2, 10), (3, 0, 11), (0, 6, 11), (0, 4, 6)],
    [(4, 11, 8), (4, 6, 11), (0, 2, 9), (2, 10, 9)],
    [(10, 9, 3), (10, 3, 2), (9, 4, 3), (11, 3, 6), (4, 6, 3)],
    [(8, 2, 3), (8, 4, 2), (4, 6, 2)],
    [(0, 4, 2), (4, 6, 2)],
    [(1, 9, 0), (2, 3, 4), (2, 4, 6), (4, 3, 8)],
    [(1, 9, 4), (1, 4, 2), (2, 4, 6)],
    [(8, 1, 3), (8, 6, 1), (8, 4, 6), (6, 10, 1)],
    [(10, 1, 0), (10, 0, 6), (6, 0, 4)],
    [(4, 6, 3), (4, 3, 8), (6, 10, 3), (0, 3, 9), (10, 9, 3)],
    [(10, 9, 4), (4, 6, 10)],
    [(4, 9, 5), (7, 6, 11)],
    [(0, 8, 3), (4, 9, 5), (11, 7, 6)],
    [(5, 0, 1), (5, 4, 0), (7, 6, 11)],
    [(11, 7, 6), (8, 3, 4), (3, 5, 4), (3, 1, 5)],
    [(9, 5, 4), (10, 1, 2), (7, 6, 11)],
    [(6, 11, 7), (1, 2, 10), (0, 8, 3), (4, 9, 5)],
    [(7, 6, 11), (5, 4, 10), (4, 2, 10), (4, 0, 2)],
    [(3, 4, 8), (3, 5, 4), (3, 2, 5), (10, 5, 2), (11, 7, 6)],
    [(7, 2, 3), (7, 6, 2), (5, 4, 9)],
    [(9, 5, 4), (0, 8, 6), (0, 6, 2), (6, 8, 7)],
    [(3, 6, 2), (3, 7, 6), (1, 5, 0), (5, 4, 0)],
    [(6, 2, 8), (6, 8, 7), (2, 1, 8), (4, 8, 5), (1, 5, 8)],
    [(9, 5, 4), (10, 1, 6), (1, 7, 6), (1, 3, 7)],
    [(1, 6, 10), (1, 7, 6), (1, 0, 7), (8, 7, 0), (9, 5, 4)],
    [(4, 0, 10), (4, 10, 5), (0, 3, 10), (6, 10, 7), (3, 7, 10)],
    [(7, 6, 10), (7, 10, 8), (5, 4, 10), (4, 8, 10)],
    [(6, 9, 5), (6, 11, 9), (11, 8, 9)],
    [(3, 6, 11), (0, 6, 3), (0, 5, 6), (0, 9, 5)],
    [(0, 11, 8), (0, 5, 11), (0, 1, 5), (5, 6, 11)],
    [(6, 11, 3), (6, 3, 5), (5, 3, 1)],
    [(1, 2, 10), (9, 5, 11), (9, 11, 8), (11, 5, 6)],
    [(0, 11, 3), (0, 6, 11), (0, 9, 6), (5, 6, 9), (1, 2, 10)],
    [(11, 8, 5), (11, 5, 6), (8, 0, 5), (10, 5, 2), (0, 2, 5)],
    [(6, 11, 3), (6, 3, 5), (2, 10, 3), (10, 5, 3)],
    [(5, 8, 9), (5, 2, 8), (5, 6, 2), (3, 8, 2)],
    [(9, 5, 6), (9, 6, 0), (0, 6, 2)],
    [(1, 5, 8), (1, 8, 0), (5, 6, 8), (3, 8, 2), (6, 2, 8)],
    [(1, 5, 6), (2, 1, 6)],
    [(1, 3, 6), (1, 6, 10), (3, 8, 6), (5, 6, 9), (8, 9, 6)],
    [(10, 1, 0), (10, 0, 6), (9, 5, 0), (5, 6, 0)],
    [(0, 3, 8), (5, 6, 10)],
    [(5, 6, 10)],
    [(11, 5, 10), (7, 5, 11)],
    [(11, 5, 10), (11, 7, 5), (8, 3, 0)],
    [(5, 11, 7), (5, 10, 11), (1, 9, 0)],
    [(10, 7, 5), (10, 11, 7), (9, 8, 1), (8, 3, 1)],
    [(11, 1, 2), (11, 7, 1), (7, 5, 1)],
    [(0, 8, 3), (1, 2, 7), (1, 7, 5), (7, 2, 11)],
    [(9, 7, 5), (9, 2, 7), (9, 0, 2), (2, 11, 7)],
    [(7, 5, 2), (7, 2, 11), (5, 9, 2), (3, 2, 8), (9, 8, 2)],
    [(2, 5, 10), (2, 3, 5), (3, 7, 5)],
    [(8, 2, 0), (8, 5, 2), (8, 7, 5), (10, 2, 5)],
    [(9, 0, 1), (5, 10, 3), (5, 3, 7), (3, 10, 2)],
    [(9, 8, 2), (9, 2, 1), (8, 7, 2), (10, 2, 5), (7, 5, 2)],
    [(1, 3, 5), (3, 7, 5)],
    [(0, 8, 7), (0, 7, 1), (1, 7, 5)],
    [(9, 0, 3), (9, 3, 5), (5, 3, 7)],
    [(9, 8, 7), (5, 9, 7)],
    [(5, 8, 4), (5, 10, 8), (10, 11, 8)],
    [(5, 0, 4), (5, 11, 0), (5, 10, 11), (11, 3, 0)],
    [(0, 1, 9), (8, 4, 10), (8, 10, 11), (10, 4, 5)],
    [(10, 11, 4), (10, 4, 5), (11, 3, 4), (9, 4, 1), (3, 1, 4)],
    [(2, 5, 1), (2, 8, 5), (2, 11, 8), (4, 5, 8)],
    [(0, 4, 11), (0, 11, 3), (4, 5, 11), (2, 11, 1), (5, 1, 11)],
    [(0, 2, 5), (0, 5, 9), (2, 11, 5), (4, 5, 8), (11, 8, 5)],
    [(9, 4, 5), (2, 11, 3)],
    [(2, 5, 10), (3, 5, 2), (3, 4, 5), (3, 8, 4)],
    [(5, 10, 2), (5, 2, 4), (4, 2, 0)],
    [(3, 10, 2), (3, 5, 10), (3, 8, 5), (4, 5, 8), (0, 1, 9)],
    [(5, 10, 2), (5, 2, 4), (1, 9, 2), (9, 4, 2)],
    [(8, 4, 5), (8, 5, 3), (3, 5, 1)],
    [(0, 4, 5), (1, 0, 5)],
    [(8, 4, 5), (8, 5, 3), (9, 0, 5), (0, 3, 5)],
    [(9, 4, 5)],
    [(4, 11, 7), (4, 9, 11), (9, 10, 11)],
    [(0, 8, 3), (4, 9, 7), (9, 11, 7), (9, 10, 11)],
    [(1, 10, 11), (1, 11, 4), (1, 4, 0), (7, 4, 11)],
    [(3, 1, 4), (3, 4, 8), (1, 10, 4), (7, 4, 11), (10, 11, 4)],
    [(4, 11, 7), (9, 11, 4), (9, 2, 11), (9, 1, 2)],
    [(9, 7, 4), (9, 11, 7), (9, 1, 11), (2, 11, 1), (0, 8, 3)],
    [(11, 7, 4), (11, 4, 2), (2, 4, 0)],
    [(11, 7, 4), (11, 4, 2), (8, 3, 4), (3, 2, 4)],
    [(2, 9, 10), (2, 7, 9), (2, 3, 7), (7, 4, 9)],
    [(9, 10, 7), (9, 7, 4), (10, 2, 7), (8, 7, 0), (2, 0, 7)],
    [(3, 7, 10), (3, 10, 2), (7, 4, 10), (1, 10, 0), (4, 0, 10)],
    [(1, 10, 2), (8, 7, 4)],
    [(4, 9, 1), (4, 1, 7), (7, 1, 3)],
    [(4, 9, 1), (4, 1, 7), (0, 8, 1), (8, 7, 1)],
    [(4, 0, 3), (7, 4, 3)],
    [(4, 8, 7)],
    [(9, 11, 8), (11, 9, 10)],
    [(3, 0, 9), (3, 9, 11), (11, 9, 10)],
    [(0, 1, 10), (0, 10, 8), (8, 10, 11)],
    [(3, 1, 11), (11, 1, 10)],
    [(1, 2, 11), (1, 11, 9), (9, 11, 8)],
    [(3, 0, 9), (3, 9, 11), (1, 2, 9), (2, 11, 9)],
    [(0, 2, 11), (11, 8, 0)],
    [(3, 2, 11)],
    [(2, 3, 8), (2, 8, 10), (10, 8, 9)],
    [(2, 0, 9), (9, 10, 2)],
    [(2, 3, 8), (2, 8, 10), (0, 1, 8), (1, 10, 8)],
    [(10, 2, 1)],
    [(1, 3, 8), (9, 1, 8)],
    [(0, 9, 1)],
    [(3, 8, 0)],
    [],
    ]

if __name__ == '__main__':
    main()
